#include <fstream>
#include <iomanip>
#include <iostream>

// glog
#include "glog/logging.h"

// gtsam
#include "gtsam/inference/Key.h"
#include "gtsam/inference/Symbol.h"
#include "gtsam/nonlinear/ISAM2.h"
#include "gtsam/nonlinear/NonlinearFactorGraph.h"
#include "gtsam/nonlinear/Values.h"

#include "gtsam/navigation/CombinedImuFactor.h"
#include "gtsam/navigation/GPSFactor.h"
#include "gtsam/navigation/ImuFactor.h"
#include "gtsam/navigation/NavState.h"

using gtsam::symbol_shorthand::B;  // Bias  (ax,ay,az,gx,gy,gz)
using gtsam::symbol_shorthand::V;  // Vel   (xdot,ydot,zdot)
using gtsam::symbol_shorthand::X;  // Pose3 (x,y,z,r,p,y)

boost::shared_ptr<gtsam::PreintegratedCombinedMeasurements::Params>
IMUParams() {
  // We use the sensor specs to build the noise model for the IMU factor.
  double accel_noise_sigma = 0.0003924;
  double gyro_noise_sigma = 0.000205689024915;
  double accel_bias_rw_sigma = 0.004905;
  double gyro_bias_rw_sigma = 0.000001454441043;
  gtsam::Matrix33 measured_acc_cov = gtsam::I_3x3 * pow(accel_noise_sigma, 2);
  gtsam::Matrix33 measured_omega_cov = gtsam::I_3x3 * pow(gyro_noise_sigma, 2);
  gtsam::Matrix33 integration_error_cov =
      gtsam::I_3x3 *
      1e-8;  // error committed in integrating position from velocities
  gtsam::Matrix33 bias_acc_cov = gtsam::I_3x3 * pow(accel_bias_rw_sigma, 2);
  gtsam::Matrix33 bias_omega_cov = gtsam::I_3x3 * pow(gyro_bias_rw_sigma, 2);
  gtsam::Matrix66 bias_acc_omega_int =
      gtsam::I_6x6 * 1e-5;  // error in the bias used for preintegration

  auto p = gtsam::PreintegratedCombinedMeasurements::Params::MakeSharedD(0.0);
  // PreintegrationBase params:
  p->accelerometerCovariance =
      measured_acc_cov;  // acc white noise in continuous
  p->integrationCovariance =
      integration_error_cov;  // integration uncertainty continuous
  // should be using 2nd order integration
  // PreintegratedRotation params:
  p->gyroscopeCovariance =
      measured_omega_cov;  // gyro white noise in continuous
  // PreintegrationCombinedMeasurements params:
  p->biasAccCovariance = bias_acc_cov;      // acc bias in continuous
  p->biasOmegaCovariance = bias_omega_cov;  // gyro bias in continuous
  p->biasAccOmegaInt = bias_acc_omega_int;

  return p;
}

int main(int argc, char** argv) {
  std::string data_file_name = argv[1];
  std::string out_file_name = "/tmp/imu_integ_fusion.txt";

  // save result to file
  std::ofstream out_fout(out_file_name);
  out_fout.precision(15);

  std::ifstream data_fin(data_file_name);
  // begin parsing the CSV file, input the first line for initization.
  // from there, we will itertate through the file and we will preintegrate
  // the IMU or add in the GPS given the input.
  gtsam::Vector10 initial_state;
  std::string value;
  std::getline(data_fin, value, ',');  // i
  for (int i = 0; i < 9; i++) {
    std::getline(data_fin, value, ',');
    initial_state(i) = std::stof(value.c_str());
  }
  std::getline(data_fin, value, '\n');
  initial_state(9) = std::stof(value.c_str());
  LOG(INFO) << "Initial states:\n" << initial_state.transpose();

  // assemble initial quaternion through GTSAM constructor
  gtsam::Rot3 prior_rotation = gtsam::Rot3::Quaternion(
      initial_state(6), initial_state(3), initial_state(4), initial_state(5));
  gtsam::Point3 prior_position(initial_state.head<3>());
  gtsam::Pose3 prior_pose(prior_rotation, prior_position);
  gtsam::Vector3 prior_velocity(initial_state.tail<3>());

  // zero initial bias
  gtsam::imuBias::ConstantBias prior_imu_bias;

  int index = 0;
  gtsam::Values initial_values;

  // insert pose at initialization
  initial_values.insert<gtsam::Pose3>(X(index), prior_pose);
  initial_values.insert<gtsam::Vector3>(V(index), prior_velocity);
  initial_values.insert<gtsam::imuBias::ConstantBias>(B(index), prior_imu_bias);

  // assemble prior noise model and add it to the graph
  auto pose_noise_model = gtsam::noiseModel::Diagonal::Sigmas(
      (gtsam::Vector(6) << 0.01, 0.01, 0.01, 0.5, 0.5, 0.5).finished());
  auto vel_noise_model = gtsam::noiseModel::Isotropic::Sigma(3, 0.1);
  auto bias_noise_model = gtsam::noiseModel::Isotropic::Sigma(6, 1e-3);

  // add all prior factors (pose, velocity, bias) to the graph
  gtsam::NonlinearFactorGraph graph;
  graph.addPrior<gtsam::Pose3>(X(index), prior_pose, pose_noise_model);
  graph.addPrior<gtsam::Vector3>(V(index), prior_velocity, vel_noise_model);
  graph.addPrior<gtsam::imuBias::ConstantBias>(
      B(index), prior_imu_bias, bias_noise_model);

  auto imu_preint_params = IMUParams();

  std::shared_ptr<gtsam::PreintegratedCombinedMeasurements> preintegrated =
      std::make_shared<gtsam::PreintegratedCombinedMeasurements>(
          imu_preint_params, prior_imu_bias);

  // store previous state for IMU integration and latest predicted outcome.
  gtsam::NavState prev_state(prior_pose, prior_velocity);
  gtsam::NavState prop_state = prev_state;
  gtsam::imuBias::ConstantBias prev_bias = prior_imu_bias;

  // iSAM2 solver
  gtsam::ISAM2Params isam2_params;
  isam2_params.relinearizeSkip = 10;
  isam2_params.relinearizeThreshold = 0.01;
  gtsam::ISAM2 isam2_solver(isam2_params);
  gtsam::Values opt_results;

  // keep track of total error over the entire run as simple performance metric.
  double curr_position_error = 0.0;
  double curr_attitude_error = 0.0;

  double output_time = 0.0;
  // the real system has noise, but here, results are nearly
  // exactly the same, so keeping this for simplicity
  double dt = 0.005;

  // all priors have been setup, now iterate through the data file.
  while (data_fin.good()) {
    std::getline(data_fin, value, ',');
    int type = std::stoi(value.c_str());
    if (type == 0) {
      // IMU measurement
      gtsam::Vector6 imu;
      for (int i = 0; i < 5; i++) {
        std::getline(data_fin, value, ',');
        imu(i) = std::stof(value.c_str());
      }
      std::getline(data_fin, value, '\n');
      imu(5) = std::stof(value.c_str());

      // adding the IMU preintegtation
      preintegrated->integrateMeasurement(imu.head<3>(), imu.tail<3>(), dt);
    } else if (type == 1) {
      gtsam::Vector7 gps;
      for (int i = 0; i < 6; ++i) {
        std::getline(data_fin, value, ',');
        gps(i) = std::stof(value.c_str());
      }
      std::getline(data_fin, value, '\n');
      gps(6) = std::stof(value.c_str());

      // adding IMU factor and GPS factor and optimizing
      index++;
      auto preint_imu_combined =
          dynamic_cast<gtsam::PreintegratedCombinedMeasurements&>(
              *preintegrated);
      // IMU factor
      gtsam::CombinedImuFactor imu_factor(X(index - 1),
                                          V(index - 1),
                                          X(index),
                                          V(index),
                                          B(index - 1),
                                          B(index),
                                          preint_imu_combined);
      graph.add(imu_factor);

      // GPS position factor
      auto correction_noise = gtsam::noiseModel::Isotropic::Sigma(3, 1.0);
      gtsam::GPSFactor gps_factor(
          X(index), gtsam::Point3(gps(0), gps(1), gps(2)), correction_noise);
      graph.add(gps_factor);

      // optimizing
      prop_state = preintegrated->predict(prev_state, prev_bias);
      initial_values.insert<gtsam::Pose3>(X(index), prop_state.pose());
      initial_values.insert<gtsam::Vector3>(V(index), prop_state.velocity());
      initial_values.insert<gtsam::imuBias::ConstantBias>(B(index), prev_bias);

      isam2_solver.update(graph, initial_values);
      opt_results = isam2_solver.calculateEstimate();

      graph.resize(0);
      initial_values.clear();

      // update previuse states
      prev_state = gtsam::NavState(opt_results.at<gtsam::Pose3>(X(index)),
                                   opt_results.at<gtsam::Vector3>(V(index)));
      prev_bias = opt_results.at<gtsam::imuBias::ConstantBias>(B(index));

      // reset the preintegration object
      preintegrated->resetIntegrationAndSetBias(prev_bias);

      // Print out the position and orientation error for comparison.
      gtsam::Vector3 result_position = prev_state.pose().translation();
      gtsam::Vector3 position_error = result_position - gps.head<3>();
      curr_position_error = position_error.norm();

      gtsam::Quaternion result_quat =
          prev_state.pose().rotation().toQuaternion();
      gtsam::Quaternion gps_quat(gps(6), gps(3), gps(4), gps(5));
      gtsam::Quaternion quat_error = result_quat * gps_quat.inverse();
      quat_error.normalize();
      gtsam::Vector3 euler_angle_error(
          quat_error.x() * 2, quat_error.y() * 2, quat_error.z() * 2);
      curr_attitude_error = euler_angle_error.norm();

      // write results to file
      out_fout << output_time << " " << prev_state.position().transpose() << " "
               << prev_state.attitude().rpy().transpose() << " "
               << prev_state.velocity().transpose() << " "
               << prev_bias.gyroscope().transpose() << " "
               << prev_bias.accelerometer().transpose() << " "
               << gps.head<3>().transpose() << std::endl;

      output_time += 1.0;
    } else {
      LOG(INFO) << "error paring file";
      return 1;
    }
  }

  // save graph to dot file
  isam2_solver.saveGraph("imu_preint_graph.dot");

  out_fout.close();
  data_fin.close();

  return 0;
}